"""
Maps of non-phase-locked ITs during
SWOT SCIENCE orbit from HYCOM

B. Yadidya
Feb 24, 2026
"""

import sys
sys.path.append('.')
from common_imports import *

dir = "data/"
ssh_ds = xr.open_dataset(dir + 'regridded_0.25deg_ALL_SSH_VR_v2.0.1.nc')
sss_ds = xr.open_dataset(dir + 'regridded_0.25deg_ALL_SLOPE_VR_v2.0.1.nc')

incoh_sd_ssh = (ssh_ds.fssh_sd_var_hycom_incoh_no_eddies)
incoh_dr_ssh = (ssh_ds.fssh_dr_var_hycom_incoh_no_eddies)
incoh_sd_ssh = incoh_sd_ssh.where(incoh_sd_ssh != 0)
incoh_dr_ssh = incoh_dr_ssh.where(incoh_dr_ssh != 0)

incoh_sd_sss_x = (sss_ds.fsss_var_x_hycom_incoh_no_eddies_sd)
incoh_sd_sss_y = (sss_ds.fsss_var_y_hycom_incoh_no_eddies_sd)
incoh_dr_sss_x = (sss_ds.fsss_var_x_hycom_incoh_no_eddies_dr)
incoh_dr_sss_y = (sss_ds.fsss_var_y_hycom_incoh_no_eddies_dr)

"""
# I used the 4km regridded datasets for the plots shown in the manuscript. 
# However for storage regions, 0.25 deg regridded datasets are uploaded on Dataverse.
ssh_ds = xr.open_dataset(dir + 'regridded_4km_ALL_SSH_VR_v2.0.1.nc')
sss_ds = xr.open_dataset(dir + 'regridded_4km_ALL_SLOPE_VR_v2.0.1.nc')

incoh_sd_ssh = coarsen_data(ssh_ds.fssh_sd_var_hycom_incoh_no_eddies, .2)
incoh_dr_ssh = coarsen_data(ssh_ds.fssh_dr_var_hycom_incoh_no_eddies, .2)
incoh_sd_ssh = incoh_sd_ssh.where(incoh_sd_ssh != 0)
incoh_dr_ssh = incoh_dr_ssh.where(incoh_dr_ssh != 0)

incoh_sd_sss_x = coarsen_data(sss_ds.fsss_var_x_hycom_incoh_no_eddies_sd, .2)
incoh_sd_sss_y = coarsen_data(sss_ds.fsss_var_y_hycom_incoh_no_eddies_sd, .2)
incoh_dr_sss_x = coarsen_data(sss_ds.fsss_var_x_hycom_incoh_no_eddies_dr, .2)
incoh_dr_sss_y = coarsen_data(sss_ds.fsss_var_y_hycom_incoh_no_eddies_dr, .2)
"""
fig_width_cm = 18.4
fig_width_in = fig_width_cm / 2.54
fig_height_max_cm = 10
fig_height_in = fig_height_max_cm / 2.54 

fig, axs = plt.subplots(
    3, 2, 
    figsize=(fig_width_in, fig_height_in), 
    constrained_layout=True
)
axs = axs.ravel()

# Define plot configurations
# (data, levels, title, unit, ticks)
plot_configs = [
    (incoh_sd_ssh, np.linspace(-1.5, 1.5, 21), 'Semidiurnal Non-phase-locked Internal Tide', 'cm$^2$', [-1.5, 0, 1.5]),
    (incoh_dr_ssh, np.linspace(-1.5, 1.5, 21), 'Diurnal Non-phase-locked Internal Tide', 'cm$^2$', [-1.5, 0, 1.5]),
    (incoh_sd_sss_x/1e-9, np.linspace(-2, 2, 21), 'Explained variance in cross-track slope', '10$^{-3}$(cm/km)$^2$', [-2, 0, 2]),
    (incoh_dr_sss_x/1e-9, np.linspace(-1, 1, 21), 'Explained variance in cross-track slope', '10$^{-3}$(cm/km)$^2$', [-1, 0, 1]),
    (incoh_sd_sss_y/1e-9, np.linspace(-2, 2, 21), 'Explained variance in along-track slope', '10$^{-3}$(cm/km)$^2$', [-2, 0, 2]),
    (incoh_dr_sss_y/1e-9, np.linspace(-1, 1, 21), 'Explained variance in along-track slope', '10$^{-3}$(cm/km)$^2$', [-1, 0, 1]),
]

# Panel labels
labels = ['A', 'B', 'C', 'D', 'E', 'F']

# Panel label style
panel_label_props = dict(
    fontsize=9, 
    fontweight='bold', 
    ha='center',
    va='center',
    bbox=dict(
        boxstyle='round, pad=0.1',
        facecolor='white',
        edgecolor='none',
        alpha=1
    )
)

for i, (data, levels, title, unit, ticks) in enumerate(plot_configs):
    ax = axs[i]
    
    # Plot contourf
    plot_obj = data.plot.contourf(
        ax=ax, levels=levels, x='longitude', cmap=colormaps.prinsenvlag_r, add_colorbar=False
    )
    
    # Set map formatting
    plot_global_map_on_axes_basemap(ax)
    ax.set_aspect('equal')
    ax.set_xlabel(''); ax.set_ylabel(''); ax.set_ylim(-60, 60)
    
    # Inset Colorbar
    cax = inset_axes(ax, width="25%", height="5%", loc='lower center',
                     bbox_to_anchor=(0.09, 0.0, 1, 1), bbox_transform=ax.transAxes)
    cbar = plt.colorbar(plot_obj, cax=cax, orientation='horizontal')
    cbar.set_ticks(ticks)
    cbar.ax.tick_params(which='minor', size=0, labelsize=7); cbar.ax.minorticks_off()
    cbar.ax.tick_params(which='major', size=0, labelsize=7)
    for label in cbar.ax.get_xticklabels(): label.set_fontweight('normal')
    cbar.ax.xaxis.set_ticks_position('top')
    cbar.ax.set_xlabel(unit, fontsize=7)
    cbar.ax.xaxis.set_label_position('top')
    
    # Panel Label
    ax.text(0.017, .92, labels[i], transform=ax.transAxes, **panel_label_props)

label_kwargs = dict(fontsize=9, ha='right', va='center', rotation=90, fontstyle = 'italic')
axs[0].text(-0.01, 0.5, 'SSH', transform=axs[0].transAxes, **label_kwargs)
axs[2].text(-0.01, 0.5, 'Cross-track Slope', transform=axs[2].transAxes, **label_kwargs)
axs[4].text(-0.01, 0.5, 'Along-track Slope', transform=axs[4].transAxes, **label_kwargs)

axs[0].set_title('Semidiurnal Non-phase-locked Internal Tide', fontsize = 9)
axs[1].set_title('Diurnal Non-phase-locked Internal Tide', fontsize = 9)
plt.savefig('Fig4_incoherent_IT_maps.png', dpi=500, bbox_inches='tight')




